{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch.nn as nn\n",
    "from tqdm import tqdm\n",
    "from typing import List, Dict\n",
    "from torch.optim import Adam\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from dataclasses import dataclass, field\n",
    "from torch_scatter import scatter_sum, scatter_softmax\n",
    "from Bio import SeqIO  # pip install biopython\n",
    "import torch_cluster\n",
    "import torch_geometric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "def calculate_rna_dihedrals(X: torch.Tensor, eps: float = 1e-7) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    计算RNA分子骨架的二面角特征并进行正弦余弦编码\n",
    "    \n",
    "    参数:\n",
    "        X: 原子坐标张量，形状为 [num_residues, num_atoms, 3]\n",
    "        eps: 数值稳定性的小常数\n",
    "        \n",
    "    返回:\n",
    "        编码后的二面角特征，形状为 [num_residues, 12]\n",
    "    \"\"\"\n",
    "    # 定义RNA骨架中的关键原子索引\n",
    "    BACKBONE_ATOMS = ['P', \"O5'\", \"C5'\", \"C4'\", \"C3'\", \"O3'\"]\n",
    "    # 每个核苷酸的二面角数量\n",
    "    DIHEDRALS_PER_NUCLEOTIDE = 6\n",
    "    \n",
    "    num_residues, num_atoms, _ = X.shape\n",
    "    \n",
    "    # 提取并重塑骨干原子坐标\n",
    "    backbone_coords = X[:, :len(BACKBONE_ATOMS), :].reshape(len(BACKBONE_ATOMS) * num_residues, 3)\n",
    "    \n",
    "    # 计算相邻原子间的单位向量\n",
    "    dX = backbone_coords[5:, :] - backbone_coords[:-5, :]\n",
    "    unit_vectors = F.normalize(dX, dim=-1)\n",
    "    \n",
    "    # 提取用于计算每个二面角的三元组向量\n",
    "    u_2 = unit_vectors[:-2, :]  # 第一个向量\n",
    "    u_1 = unit_vectors[1:-1, :] # 中间向量\n",
    "    u_0 = unit_vectors[2:, :]   # 最后一个向量\n",
    "    \n",
    "    # 计算相邻平面的法向量\n",
    "    n_2 = F.normalize(torch.cross(u_2, u_1), dim=-1)\n",
    "    n_1 = F.normalize(torch.cross(u_1, u_0), dim=-1)\n",
    "    \n",
    "    # 计算二面角的余弦值\n",
    "    cosD = (n_2 * n_1).sum(-1)\n",
    "    cosD = torch.clamp(cosD, -1 + eps, 1 - eps)\n",
    "    \n",
    "    # 计算二面角，包括符号确定\n",
    "    angle_magnitude = torch.acos(cosD)\n",
    "    angle_sign = torch.sign((u_2 * n_1).sum(-1))\n",
    "    dihedral_angles = angle_sign * angle_magnitude\n",
    "    \n",
    "    # 填充并重塑二面角张量以匹配原始序列长度\n",
    "    padded_angles = F.pad(dihedral_angles, (3, 4), 'constant', 0)\n",
    "    reshaped_angles = padded_angles.view(num_residues, DIHEDRALS_PER_NUCLEOTIDE)\n",
    "    \n",
    "    # 使用正弦和余弦编码二面角，增强特征表达能力\n",
    "    return torch.cat((torch.cos(reshaped_angles), torch.sin(reshaped_angles)), dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "def normalize_vectors(vectors: torch.Tensor, dim: int = -1) -> torch.Tensor:\n",
    "    \"\"\"归一化向量，处理零向量以提高数值稳定性\"\"\"\n",
    "    return F.normalize(vectors + 1e-12, dim=dim)\n",
    "\n",
    "def compute_local_frames(X: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n",
    "    \"\"\"\n",
    "    为每个核苷酸计算局部坐标系（以C3'原子为中心）\n",
    "    \n",
    "    返回:\n",
    "        X_c3: C3'原子坐标 [B, N, 3]\n",
    "        Q: 局部坐标系旋转矩阵 [B, N, 3, 3]\n",
    "    \"\"\"\n",
    "    # 提取C3'原子坐标 (索引4)\n",
    "    X_c3 = X[:, :, 4, :]  # [B, N, 3]\n",
    "    \n",
    "    # 计算相邻C3'原子间的单位向量\n",
    "    dX = X_c3[:, 1:, :] - X_c3[:, :-1, :]\n",
    "    u = normalize_vectors(dX)\n",
    "    \n",
    "    # 构建局部坐标系\n",
    "    # b1: 沿链方向的向量\n",
    "    b1 = F.pad(u, (0, 0, 0, 1), 'replicate')  # 填充最后一个位置\n",
    "    \n",
    "    # n0: 垂直于链的法向量\n",
    "    # 使用相邻两个向量的叉乘\n",
    "    u_prev = F.pad(u[:, :-1, :], (0, 0, 1, 0), 'replicate')\n",
    "    n0 = normalize_vectors(torch.cross(u_prev, u))\n",
    "    n0 = F.pad(n0, (0, 0, 0, 1), 'replicate')  # 填充最后一个位置\n",
    "    \n",
    "    # b3: b1和n0的叉乘，完成正交基\n",
    "    b3 = normalize_vectors(torch.cross(b1, n0))\n",
    "    \n",
    "    # 构建旋转矩阵 [B, N, 3, 3]\n",
    "    Q = torch.stack([b1, n0, b3], dim=2)\n",
    "    \n",
    "    return X_c3, Q\n",
    "\n",
    "def calculate_v_direct(X: torch.Tensor, Q: torch.Tensor, X_c3: torch.Tensor) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    计算内部方向特征 (V_direct)\n",
    "    \n",
    "    参数:\n",
    "        X: 原子坐标 [B, N, A, 3]\n",
    "        Q: 局部坐标系旋转矩阵 [B, N, 3, 3]\n",
    "        X_c3: C3'原子坐标 [B, N, 3]\n",
    "        \n",
    "    返回:\n",
    "        V_direct: 内部方向特征 [B, N, 9]\n",
    "    \"\"\"\n",
    "    B, N, _, _ = X.shape\n",
    "    \n",
    "    # 选择内部关键原子: P(0), C5'(2), C4'(3)\n",
    "    atoms_of_interest = [0, 2, 3]\n",
    "    X_inner = X[:, :, atoms_of_interest, :]  # [B, N, 3, 3]\n",
    "    \n",
    "    # 计算相对于C3'的位移向量\n",
    "    dX_inner = X_inner - X_c3.unsqueeze(2)  # [B, N, 3, 3]\n",
    "    \n",
    "    # 转换到局部坐标系\n",
    "    dU_inner = torch.matmul(Q.unsqueeze(2), dX_inner.unsqueeze(-1)).squeeze(-1)  # [B, N, 3, 3]\n",
    "    \n",
    "    # 归一化方向向量\n",
    "    V_direct = normalize_vectors(dU_inner.reshape(B, N, -1))  # [B, N, 9]\n",
    "    \n",
    "    return V_direct\n",
    "\n",
    "def calculate_e_direct(X: torch.Tensor, Q: torch.Tensor, X_c3: torch.Tensor, E_idx: torch.Tensor) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    计算邻居方向特征 (E_direct)\n",
    "    \n",
    "    参数:\n",
    "        X: 原子坐标 [B, N, A, 3]\n",
    "        Q: 局部坐标系旋转矩阵 [B, N, 3, 3]\n",
    "        X_c3: C3'原子坐标 [B, N, 3]\n",
    "        E_idx: 邻居索引 [B, N, K]\n",
    "        \n",
    "    返回:\n",
    "        E_direct: 邻居方向特征 [B, N, K, 3]\n",
    "    \"\"\"\n",
    "    B, N, K = E_idx.shape\n",
    "    \n",
    "    # 收集邻居的关键原子坐标\n",
    "    key_atoms = [0, 1, 2, 3, 5]  # P, O5', C5', C4', O3'\n",
    "    X_neighbors = torch.stack([\n",
    "        gather_nodes(X[:, :, atom_idx, :], E_idx) \n",
    "        for atom_idx in key_atoms\n",
    "    ], dim=3)  # [B, N, K, 5, 3]\n",
    "    \n",
    "    # 计算相对于中心C3'的位移向量\n",
    "    dX = X_neighbors - X_c3.unsqueeze(2).unsqueeze(3)  # [B, N, K, 5, 3]\n",
    "    \n",
    "    # 转换到局部坐标系\n",
    "    Q_expanded = Q.unsqueeze(2).unsqueeze(3)  # [B, N, 1, 1, 3, 3]\n",
    "    dU = torch.matmul(Q_expanded, dX.unsqueeze(-1)).squeeze(-1)  # [B, N, K, 5, 3]\n",
    "    \n",
    "    # 合并多个原子的方向特征\n",
    "    E_direct = normalize_vectors(dU.mean(dim=3))  # [B, N, K, 3]\n",
    "    \n",
    "    return E_direct\n",
    "\n",
    "def calculate_e_orient(Q: torch.Tensor, E_idx: torch.Tensor) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    计算方向关系特征 (E_orient)\n",
    "    \n",
    "    参数:\n",
    "        Q: 局部坐标系旋转矩阵 [B, N, 3, 3]\n",
    "        E_idx: 邻居索引 [B, N, K]\n",
    "        \n",
    "    返回:\n",
    "        E_orient: 方向关系特征 [B, N, K, 4] (四元数表示)\n",
    "    \"\"\"\n",
    "    B, N, K = E_idx.shape\n",
    "    \n",
    "    # 收集邻居的局部坐标系\n",
    "    Q_neighbors = gather_nodes(Q.reshape(B, N, -1), E_idx).reshape(B, N, K, 3, 3)  # [B, N, K, 3, 3]\n",
    "    \n",
    "    # 计算旋转矩阵: 当前坐标系 -> 邻居坐标系\n",
    "    R = torch.matmul(Q.transpose(-1, -2).unsqueeze(2), Q_neighbors)  # [B, N, K, 3, 3]\n",
    "    \n",
    "    # 将旋转矩阵转换为四元数表示\n",
    "    E_orient = rotation_matrix_to_quaternion(R)  # [B, N, K, 4]\n",
    "    \n",
    "    return E_orient\n",
    "\n",
    "def rotation_matrix_to_quaternion(R: torch.Tensor) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    将旋转矩阵转换为四元数表示\n",
    "    \n",
    "    参数:\n",
    "        R: 旋转矩阵 [B, N, K, 3, 3]\n",
    "        \n",
    "    返回:\n",
    "        quaternion: 四元数 [B, N, K, 4]\n",
    "    \"\"\"\n",
    "    batch_dims = R.shape[:-2]\n",
    "    R = R.reshape(-1, 3, 3)\n",
    "    \n",
    "    trace = R[:, 0, 0] + R[:, 1, 1] + R[:, 2, 2]\n",
    "    q = torch.zeros((R.shape[0], 4), device=R.device)\n",
    "    \n",
    "    mask1 = trace > 0\n",
    "    if mask1.any():\n",
    "        s = torch.sqrt(trace[mask1] + 1.0) * 2\n",
    "        q[mask1, 0] = 0.25 * s\n",
    "        q[mask1, 1] = (R[mask1, 2, 1] - R[mask1, 1, 2]) / s\n",
    "        q[mask1, 2] = (R[mask1, 0, 2] - R[mask1, 2, 0]) / s\n",
    "        q[mask1, 3] = (R[mask1, 1, 0] - R[mask1, 0, 1]) / s\n",
    "    \n",
    "    mask2 = (trace <= 0) & (R[:, 0, 0] > R[:, 1, 1]) & (R[:, 0, 0] > R[:, 2, 2])\n",
    "    if mask2.any():\n",
    "        s = torch.sqrt(1.0 + R[mask2, 0, 0] - R[mask2, 1, 1] - R[mask2, 2, 2]) * 2\n",
    "        q[mask2, 0] = (R[mask2, 2, 1] - R[mask2, 1, 2]) / s\n",
    "        q[mask2, 1] = 0.25 * s\n",
    "        q[mask2, 2] = (R[mask2, 0, 1] + R[mask2, 1, 0]) / s\n",
    "        q[mask2, 3] = (R[mask2, 0, 2] + R[mask2, 2, 0]) / s\n",
    "    \n",
    "    mask3 = (trace <= 0) & (R[:, 1, 1] > R[:, 2, 2])\n",
    "    if mask3.any():\n",
    "        s = torch.sqrt(1.0 + R[mask3, 1, 1] - R[mask3, 0, 0] - R[mask3, 2, 2]) * 2\n",
    "        q[mask3, 0] = (R[mask3, 0, 2] - R[mask3, 2, 0]) / s\n",
    "        q[mask3, 1] = (R[mask3, 0, 1] + R[mask3, 1, 0]) / s\n",
    "        q[mask3, 2] = 0.25 * s\n",
    "        q[mask3, 3] = (R[mask3, 1, 2] + R[mask3, 2, 1]) / s\n",
    "    \n",
    "    mask4 = (trace <= 0) & (R[:, 2, 2] > R[:, 0, 0]) & (R[:, 2, 2] > R[:, 1, 1])\n",
    "    if mask4.any():\n",
    "        s = torch.sqrt(1.0 + R[mask4, 2, 2] - R[mask4, 0, 0] - R[mask4, 1, 1]) * 2\n",
    "        q[mask4, 0] = (R[mask4, 1, 0] - R[mask4, 0, 1]) / s\n",
    "        q[mask4, 1] = (R[mask4, 0, 2] + R[mask4, 2, 0]) / s\n",
    "        q[mask4, 2] = (R[mask4, 1, 2] + R[mask4, 2, 1]) / s\n",
    "        q[mask4, 3] = 0.25 * s\n",
    "    \n",
    "    # 归一化四元数\n",
    "    q = F.normalize(q, dim=-1)\n",
    "    \n",
    "    return q.reshape(*batch_dims, 4)\n",
    "\n",
    "def gather_nodes(nodes: torch.Tensor, neighbor_idx: torch.Tensor) -> torch.Tensor:\n",
    "    \"\"\"从节点特征中收集邻居特征\"\"\"\n",
    "    batch_size, num_nodes = neighbor_idx.shape[:2]\n",
    "    neighbors_flat = neighbor_idx.view(batch_size, -1)\n",
    "    neighbors_flat = neighbors_flat.unsqueeze(-1).expand(-1, -1, nodes.size(2))\n",
    "    neighbor_features = torch.gather(nodes, 1, neighbors_flat)\n",
    "    return neighbor_features.view(batch_size, num_nodes, neighbor_idx.shape[2], -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNAFeatures(nn.Module):\n",
    "    def __init__(self, edge_features, node_features, node_feat_types=[], edge_feat_types=[], num_rbf=16, top_k=30, augment_eps=0., dropout=0.1, args=None):\n",
    "        super(RNAFeatures, self).__init__()\n",
    "        \"\"\"Extract RNA Features\"\"\"\n",
    "        self.edge_features = edge_features\n",
    "        self.node_features = node_features\n",
    "        self.top_k = top_k\n",
    "        self.augment_eps = augment_eps \n",
    "        self.num_rbf = num_rbf\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.node_feat_types = node_feat_types\n",
    "        self.edge_feat_types = edge_feat_types\n",
    "\n",
    "        node_in = sum([feat_dims['node'][feat] for feat in node_feat_types])\n",
    "        edge_in = sum([feat_dims['edge'][feat] for feat in edge_feat_types])\n",
    "        self.node_embedding = nn.Linear(node_in,  node_features, bias=True)\n",
    "        self.edge_embedding = nn.Linear(edge_in, edge_features, bias=True)\n",
    "        self.norm_nodes = Normalize(node_features)\n",
    "        self.norm_edges = Normalize(edge_features)\n",
    "        \n",
    "    def _dist(self, X, mask, eps=1E-6):\n",
    "        mask_2D = torch.unsqueeze(mask,1) * torch.unsqueeze(mask,2)\n",
    "        dX = torch.unsqueeze(X,1) - torch.unsqueeze(X,2)\n",
    "        D = (1. - mask_2D)*10000 + mask_2D* torch.sqrt(torch.sum(dX**2, 3) + eps)\n",
    "\n",
    "        D_max, _ = torch.max(D, -1, keepdim=True)\n",
    "        D_adjust = D + (1. - mask_2D) * (D_max+1)\n",
    "        D_neighbors, E_idx = torch.topk(D_adjust, min(self.top_k, D_adjust.shape[-1]), dim=-1, largest=False)\n",
    "        return D_neighbors, E_idx\n",
    "    \n",
    "    def _rbf(self, D):\n",
    "        D_min, D_max, D_count = 0., 20., self.num_rbf\n",
    "        D_mu = torch.linspace(D_min, D_max, D_count, device=D.device)\n",
    "        D_mu = D_mu.view([1,1,1,-1])\n",
    "        D_sigma = (D_max - D_min) / D_count\n",
    "        D_expand = torch.unsqueeze(D, -1)\n",
    "        return torch.exp(-((D_expand - D_mu) / D_sigma)**2)\n",
    "    \n",
    "    def _get_rbf(self, A, B, E_idx=None, num_rbf=16):\n",
    "        if E_idx is not None:\n",
    "            D_A_B = torch.sqrt(torch.sum((A[:,:,None,:] - B[:,None,:,:])**2,-1) + 1e-6)\n",
    "            D_A_B_neighbors = gather_edges(D_A_B[:,:,:,None], E_idx)[:,:,:,0]\n",
    "            RBF_A_B = self._rbf(D_A_B_neighbors)\n",
    "        else:\n",
    "            D_A_B = torch.sqrt(torch.sum((A[:,:,None,:] - B[:,:,None,:])**2,-1) + 1e-6)\n",
    "            RBF_A_B = self._rbf(D_A_B)\n",
    "        return RBF_A_B\n",
    "    \n",
    "    def _quaternions(self, R):\n",
    "        diag = torch.diagonal(R, dim1=-2, dim2=-1)\n",
    "        Rxx, Ryy, Rzz = diag.unbind(-1)\n",
    "        magnitudes = 0.5 * torch.sqrt(torch.abs(1 + torch.stack([\n",
    "              Rxx - Ryy - Rzz, \n",
    "            - Rxx + Ryy - Rzz, \n",
    "            - Rxx - Ryy + Rzz\n",
    "        ], -1)))\n",
    "        _R = lambda i,j: R[:,:,:,i,j]\n",
    "        signs = torch.sign(torch.stack([\n",
    "            _R(2,1) - _R(1,2),\n",
    "            _R(0,2) - _R(2,0),\n",
    "            _R(1,0) - _R(0,1)\n",
    "        ], -1))\n",
    "        xyz = signs * magnitudes\n",
    "        w = torch.sqrt(F.relu(1 + diag.sum(-1, keepdim=True))) / 2.\n",
    "        Q = torch.cat((xyz, w), -1)\n",
    "        Q = F.normalize(Q, dim=-1)\n",
    "        return Q\n",
    "    \n",
    "    def _orientations_coarse(self, X, E_idx, eps=1e-6):\n",
    "        V = X.clone()\n",
    "        X = X[:,:,:6,:].reshape(X.shape[0], 6*X.shape[1], 3) \n",
    "        dX = X[:,1:,:] - X[:,:-1,:]\n",
    "        U = _normalize(dX, dim=-1)\n",
    "        u_0, u_1 = U[:,:-2,:], U[:,1:-1,:]\n",
    "        n_0 = _normalize(torch.cross(u_0, u_1), dim=-1)\n",
    "        b_1 = _normalize(u_0 - u_1, dim=-1)\n",
    "        \n",
    "        # select C3'\n",
    "        n_0 = n_0[:,4::6,:] \n",
    "        b_1 = b_1[:,4::6,:]\n",
    "        X = X[:,4::6,:]\n",
    "\n",
    "        Q = torch.stack((b_1, n_0, torch.cross(b_1, n_0)), 2)\n",
    "        Q = Q.view(list(Q.shape[:2]) + [9])\n",
    "        Q = F.pad(Q, (0,0,0,1), 'constant', 0) # [16, 464, 9]\n",
    "\n",
    "        Q_neighbors = gather_nodes(Q, E_idx) # [16, 464, 30, 9]\n",
    "        P_neighbors = gather_nodes(V[:,:,0,:], E_idx) # [16, 464, 30, 3]\n",
    "        O5_neighbors = gather_nodes(V[:,:,1,:], E_idx)\n",
    "        C5_neighbors = gather_nodes(V[:,:,2,:], E_idx)\n",
    "        C4_neighbors = gather_nodes(V[:,:,3,:], E_idx)\n",
    "        O3_neighbors = gather_nodes(V[:,:,5,:], E_idx)\n",
    "        \n",
    "        Q = Q.view(list(Q.shape[:2]) + [3,3]).unsqueeze(2) # [16, 464, 1, 3, 3]\n",
    "        Q_neighbors = Q_neighbors.view(list(Q_neighbors.shape[:3]) + [3,3]) # [16, 464, 30, 3, 3]\n",
    "\n",
    "        dX = torch.stack([P_neighbors,O5_neighbors,C5_neighbors,C4_neighbors,O3_neighbors], dim=3) - X[:,:,None,None,:] # [16, 464, 30, 3]\n",
    "        dU = torch.matmul(Q[:,:,:,None,:,:], dX[...,None]).squeeze(-1) # [16, 464, 30, 3] 邻居的相对坐标\n",
    "        B, N, K = dU.shape[:3]\n",
    "        E_direct = _normalize(dU, dim=-1)\n",
    "        E_direct = E_direct.reshape(B, N, K,-1)\n",
    "        R = torch.matmul(Q.transpose(-1,-2), Q_neighbors)\n",
    "        E_orient = self._quaternions(R)\n",
    "        \n",
    "        dX_inner = V[:,:,[0,2,3],:] - X.unsqueeze(-2)\n",
    "        dU_inner = torch.matmul(Q, dX_inner.unsqueeze(-1)).squeeze(-1)\n",
    "        dU_inner = _normalize(dU_inner, dim=-1)\n",
    "        V_direct = dU_inner.reshape(B,N,-1)\n",
    "        return V_direct, E_direct, E_orient\n",
    "    \n",
    "    def _dihedrals(self, X, eps=1e-7):\n",
    "        # P, O5', C5', C4', C3', O3'\n",
    "        X = X[:,:,:6,:].reshape(X.shape[0], 6*X.shape[1], 3)\n",
    "\n",
    "        # Shifted slices of unit vectors\n",
    "        # https://iupac.qmul.ac.uk/misc/pnuc2.html#220\n",
    "        # https://x3dna.org/highlights/torsion-angles-of-nucleic-acid-structures\n",
    "        # alpha:   O3'_{i-1} P_i O5'_i C5'_i\n",
    "        # beta:    P_i O5'_i C5'_i C4'_i\n",
    "        # gamma:   O5'_i C5'_i C4'_i C3'_i\n",
    "        # delta:   C5'_i C4'_i C3'_i O3'_i\n",
    "        # epsilon: C4'_i C3'_i O3'_i P_{i+1}\n",
    "        # zeta:    C3'_i O3'_i P_{i+1} O5'_{i+1} \n",
    "        # What's more:\n",
    "        #   chi: C1' - N9 \n",
    "        #   chi is different for (C, T, U) and (A, G) https://x3dna.org/highlights/the-chi-x-torsion-angle-characterizes-base-sugar-relative-orientation\n",
    "\n",
    "        dX = X[:, 5:, :] - X[:, :-5, :] # O3'-P, P-O5', O5'-C5', C5'-C4', ...\n",
    "        U = F.normalize(dX, dim=-1)\n",
    "        u_2 = U[:,:-2,:]  # O3'-P, P-O5', ...\n",
    "        u_1 = U[:,1:-1,:] # P-O5', O5'-C5', ...\n",
    "        u_0 = U[:,2:,:]   # O5'-C5', C5'-C4', ...\n",
    "        # Backbone normals\n",
    "        n_2 = F.normalize(torch.cross(u_2, u_1), dim=-1)\n",
    "        n_1 = F.normalize(torch.cross(u_1, u_0), dim=-1)\n",
    "\n",
    "        # Angle between normals\n",
    "        cosD = (n_2 * n_1).sum(-1)\n",
    "        cosD = torch.clamp(cosD, -1+eps, 1-eps)\n",
    "        D = torch.sign((u_2 * n_1).sum(-1)) * torch.acos(cosD)\n",
    "        \n",
    "        D = F.pad(D, (3,4), 'constant', 0)\n",
    "        D = D.view((D.size(0), D.size(1) //6, 6))\n",
    "        return torch.cat((torch.cos(D), torch.sin(D)), 2) # return D_features\n",
    "    \n",
    "    def forward(self, X, S, mask):\n",
    "        if self.training and self.augment_eps > 0:\n",
    "            X = X + self.augment_eps * torch.randn_like(X)\n",
    "\n",
    "        # Build k-Nearest Neighbors graph\n",
    "        B, N, _,_ = X.shape\n",
    "        # P, O5', C5', C4', C3', O3'\n",
    "        atom_P = X[:, :, 0, :]\n",
    "        atom_O5_ = X[:, :, 1, :]\n",
    "        atom_C5_ = X[:, :, 2, :]\n",
    "        atom_C4_ = X[:, :, 3, :]\n",
    "        atom_C3_ = X[:, :, 4, :] \n",
    "        atom_O3_ = X[:, :, 5, :]\n",
    "\n",
    "        X_backbone = atom_P\n",
    "        D_neighbors, E_idx = self._dist(X_backbone, mask)        \n",
    "\n",
    "        mask_bool = (mask==1)\n",
    "        mask_attend = gather_nodes(mask.unsqueeze(-1), E_idx).squeeze(-1)\n",
    "        mask_attend = (mask.unsqueeze(-1) * mask_attend) == 1\n",
    "        edge_mask_select = lambda x: torch.masked_select(x, mask_attend.unsqueeze(-1)).reshape(-1,x.shape[-1])\n",
    "        node_mask_select = lambda x: torch.masked_select(x, mask_bool.unsqueeze(-1)).reshape(-1, x.shape[-1])\n",
    "\n",
    "        # node features\n",
    "        h_V = []\n",
    "        # angle\n",
    "        V_angle = node_mask_select(self._dihedrals(X))  # 计算二面角\n",
    "        # distance\n",
    "        node_list = ['O5_-P', 'C5_-P', 'C4_-P', 'C3_-P', 'O3_-P']\n",
    "        V_dist = []\n",
    "        # 原子内部距离编码\n",
    "        for pair in node_list:\n",
    "            atom1, atom2 = pair.split('-')\n",
    "            V_dist.append(node_mask_select(self._get_rbf(vars()['atom_' + atom1], vars()['atom_' + atom2], None, self.num_rbf).squeeze()))\n",
    "        V_dist = torch.cat(tuple(V_dist), dim=-1).squeeze()\n",
    "        # direction\n",
    "        # V_direct 内部原子的方向，E_direct 邻居原子间的方向 邻居原子间的旋转向量E_orient\n",
    "        V_direct, E_direct, E_orient = self._orientations_coarse(X, E_idx)\n",
    "        V_direct = node_mask_select(V_direct)\n",
    "        E_direct, E_orient = list(map(lambda x: edge_mask_select(x), [E_direct, E_orient]))\n",
    "\n",
    "        # edge features\n",
    "        h_E = []\n",
    "        # dist\n",
    "        edge_list = ['P-P', 'O5_-P', 'C5_-P', 'C4_-P', 'C3_-P', 'O3_-P']\n",
    "        E_dist = [] \n",
    "        for pair in edge_list:\n",
    "            atom1, atom2 = pair.split('-')\n",
    "            E_dist.append(edge_mask_select(self._get_rbf(vars()['atom_' + atom1], vars()['atom_' + atom2], E_idx, self.num_rbf)))\n",
    "        E_dist = torch.cat(tuple(E_dist), dim=-1)\n",
    "\n",
    "        if 'angle' in self.node_feat_types:\n",
    "            h_V.append(V_angle)\n",
    "        if 'distance' in self.node_feat_types:\n",
    "            h_V.append(V_dist)\n",
    "        if 'direction' in self.node_feat_types:\n",
    "            h_V.append(V_direct)\n",
    "\n",
    "        if 'orientation' in self.edge_feat_types:\n",
    "            h_E.append(E_orient)\n",
    "        if 'distance' in self.edge_feat_types:\n",
    "            h_E.append(E_dist)\n",
    "        if 'direction' in self.edge_feat_types:\n",
    "            h_E.append(E_direct)\n",
    "\n",
    "        return X, S, h_V, h_E, E_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch.nn as nn\n",
    "from tqdm import tqdm\n",
    "from typing import List, Dict\n",
    "from torch.optim import Adam\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from dataclasses import dataclass, field\n",
    "from torch_scatter import scatter_sum, scatter_softmax\n",
    "from Bio import SeqIO  # pip install biopython\n",
    "import torch_cluster\n",
    "import torch_geometric\n",
    "\n",
    "\n",
    "NUM_TO_LETTER = np.array(['A', 'G', 'C', 'U'])\n",
    "LETTER_TO_NUM = {'A': 0, 'G': 1, 'C': 2, 'U': 3}\n",
    "\n",
    "def get_posenc(edge_index, num_posenc=16):\n",
    "    # From https://github.com/jingraham/neurips19-graph-protein-design\n",
    "    num_posenc = num_posenc\n",
    "    d = edge_index[0] - edge_index[1]\n",
    "\n",
    "    frequency = torch.exp(\n",
    "        torch.arange(0, num_posenc, 2, dtype=torch.float32, device=d.device)\n",
    "        * -(np.log(10000.0) / num_posenc)\n",
    "    )\n",
    "\n",
    "    angles = d.unsqueeze(-1) * frequency\n",
    "    E = torch.cat((torch.cos(angles), torch.sin(angles)), -1)\n",
    "    return E\n",
    "\n",
    "\n",
    "def get_orientations(X):\n",
    "    # X : num_conf x num_res x 3\n",
    "    forward = normalize(X[:, 1:] - X[:, :-1])\n",
    "    backward = normalize(X[:, :-1] - X[:, 1:])\n",
    "    forward = F.pad(forward, [0, 0, 0, 1])\n",
    "    backward = F.pad(backward, [0, 0, 1, 0])\n",
    "    return torch.cat([forward.unsqueeze(-2), backward.unsqueeze(-2)], -2)\n",
    "\n",
    "\n",
    "def get_orientations_single(X):\n",
    "    # X : num_res x 3\n",
    "    forward = normalize(X[1:] - X[:-1])\n",
    "    backward = normalize(X[:-1] - X[1:])\n",
    "    forward = F.pad(forward, [0, 0, 0, 1])\n",
    "    backward = F.pad(backward, [0, 0, 1, 0])\n",
    "    return torch.cat([forward.unsqueeze(-2), backward.unsqueeze(-2)], -2)\n",
    "\n",
    "def get_sidechains(X):\n",
    "    # X : num_conf x num_res x 3 x 3\n",
    "    p, origin, n = X[:, :, 0], X[:, :, 1], X[:, :, 2]\n",
    "    n, p = normalize(n - origin), normalize(p - origin)\n",
    "    return torch.cat([n.unsqueeze_(-2), p.unsqueeze_(-2)], -2)\n",
    "\n",
    "def get_sidechains_single(X):\n",
    "    # X : num_res x 3 x 3\n",
    "    p, origin, n = X[:, 0], X[:, 1], X[:, 2]\n",
    "    n, p = normalize(n - origin), normalize(p - origin)\n",
    "    return torch.cat([n.unsqueeze_(-2), p.unsqueeze_(-2)], -2)\n",
    "\n",
    "def normalize(tensor, dim=-1):\n",
    "    '''\n",
    "    Normalizes a `torch.Tensor` along dimension `dim` without `nan`s.\n",
    "    '''\n",
    "    return torch.nan_to_num(\n",
    "        torch.div(tensor, torch.linalg.norm(tensor, dim=dim, keepdim=True)))\n",
    "\n",
    "\n",
    "def rbf(D, D_min=0., D_max=20., D_count=16):\n",
    "    '''\n",
    "    From https://github.com/jingraham/neurips19-graph-protein-design\n",
    "\n",
    "    Returns an RBF embedding of `torch.Tensor` `D` along a new axis=-1.\n",
    "    That is, if `D` has shape [...dims], then the returned tensor will have\n",
    "    shape [...dims, D_count].\n",
    "\n",
    "    TODO switch to DimeNet RBFs\n",
    "    '''\n",
    "    D_mu = torch.linspace(D_min, D_max, D_count, device=D.device)  # 起点 终点 步数\n",
    "    D_mu = D_mu.view([1, -1])\n",
    "    D_sigma = (D_max - D_min) / D_count\n",
    "    D_expand = torch.unsqueeze(D, -1)\n",
    "\n",
    "    RBF = torch.exp(-((D_expand - D_mu) / D_sigma) ** 2)\n",
    "    return RBF\n",
    "\n",
    "@torch.no_grad()\n",
    "def construct_data_single(coords, seq=None, mask=None, num_posenc=16, num_rbf=16, knn_num=10):\n",
    "    \"\"\"\n",
    "    使用C1' 作为特征点进行聚类\n",
    "        目的: 使用 C1' 作为节点位置，因其能较好地反映糖环的构象。\n",
    "    \"\"\"\n",
    "    coords = torch.as_tensor(coords, dtype=torch.float32) # num_res x 3 x 3\n",
    "    # seq is np.array/string, convert to torch.tensor\n",
    "    if isinstance(seq, np.ndarray):\n",
    "        seq = torch.as_tensor(seq, dtype=torch.long)\n",
    "    else:\n",
    "        seq = torch.as_tensor(\n",
    "            [LETTER_TO_NUM[residue] for residue in seq],\n",
    "            dtype=torch.long\n",
    "        )\n",
    "\n",
    "    # Compute features\n",
    "    # node positions: num_res x 3\n",
    "    coord_C = coords[:, 1].clone()\n",
    "    # Construct merged edge index\n",
    "    edge_index = torch_cluster.knn_graph(coord_C, k=knn_num)  # k近邻建图\n",
    "    edge_index = torch_geometric.utils.coalesce(edge_index)  # 边去重，没属性就只去重，有属性，重复的边属性特征按参数聚合\n",
    "\n",
    "    # Node attributes: num_res x 2 x 3, each\n",
    "    orientations = get_orientations_single(coord_C)  # 节点间的相对坐标\n",
    "    sidechains = get_sidechains_single(coords)  # 节点内的相对坐标\n",
    "\n",
    "    # Edge displacement vectors: num_edges x  3\n",
    "    edge_vectors = coord_C[edge_index[0]] - coord_C[edge_index[1]]\n",
    "\n",
    "    # Edge RBF features: num_edges x num_rbf\n",
    "    edge_rbf = rbf(edge_vectors.norm(dim=-1), D_count=num_rbf)  # 编码相邻节点的物理距离，径向基函数，距离转特征的常见方法\n",
    "    # Edge positional encodings: num_edges x num_posenc\n",
    "    edge_posenc = get_posenc(edge_index, num_posenc)  # 编码相邻节点间的索引距离\n",
    "\n",
    "    node_s = (seq.unsqueeze(-1) == torch.arange(4).unsqueeze(0)).float()  # 节点标签\n",
    "    node_v = torch.cat([orientations, sidechains], dim=-2)  # 节点特征\n",
    "    edge_s = torch.cat([edge_rbf, edge_posenc], dim=-1)  # 每条边的特征\n",
    "    edge_v = normalize(edge_vectors).unsqueeze(-2)  # 每条边的连接向量\n",
    "\n",
    "    node_s, node_v, edge_s, edge_v = map(\n",
    "        torch.nan_to_num,\n",
    "        (node_s, node_v, edge_s, edge_v)\n",
    "    )\n",
    "\n",
    "    # add mask for invalid residues\n",
    "    if mask is None:\n",
    "        mask = coords.sum(dim=(1, 2)) != 0.\n",
    "    else:\n",
    "        mask = torch.tensor(mask)\n",
    "    \n",
    "    return {\n",
    "        'seq': seq,  # [L], 碱基类型的数字表示\n",
    "        'coords': coords,  # [L, 3, 3], 残基坐标\n",
    "        'node_s': node_s,  # [L, 4], 碱基独热编码\n",
    "        'node_v': node_v,  # [L, 3, 3], 节点方向/侧链特征\n",
    "        'edge_s': edge_s,  # [E, num_rbf+num_posenc], 边的标量特征\n",
    "        'edge_v': edge_v,  # [E, 1, 3], 边的向量特征（方向）\n",
    "        'edge_index': edge_index,  # [2, E], 图的边索引\n",
    "        'mask': mask  # [L], 有效残基掩码\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNADatasetV2(Dataset):\n",
    "    def __init__(self, data_path, is_train=True):\n",
    "        super(RNADatasetV2, self).__init__()\n",
    "        self.npy_dir = data_path + \"/coords\"\n",
    "        self.name_list = [i[:-4] for i in os.listdir(data_path + \"/coords\")]\n",
    "        self.seq_dir =  data_path + \"/seqs/\"\n",
    "        self.cache = {}\n",
    "        self.is_train = is_train\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.name_list)\n",
    "\n",
    "    def get_pdb_id(self, idx):\n",
    "        return self.name_list[idx]\n",
    "\n",
    "    def merge_coords_file_path(self, pdb_id):\n",
    "        return os.path.join(self.npy_dir, pdb_id + '.npy')\n",
    "\n",
    "    def load_feature(self, pdb_id):\n",
    "        coords = np.load(self.merge_coords_file_path(pdb_id))\n",
    "        feature = {\n",
    "            \"name\": pdb_id,\n",
    "            \"coords\": {\n",
    "                \"P\": coords[:, 0, :],\n",
    "                \"O5'\": coords[:, 1, :],\n",
    "                \"C5'\": coords[:, 2, :],\n",
    "                \"C4'\": coords[:, 3, :],\n",
    "                \"C3'\": coords[:, 4, :],\n",
    "                \"O3'\": coords[:, 5, :],\n",
    "                \"N\": coords[:, 6, :],\n",
    "            }\n",
    "        }\n",
    "\n",
    "        return feature\n",
    "    \n",
    "    def read_fasta_biopython(self, file_path):\n",
    "        sequences = {}\n",
    "        for record in SeqIO.parse(file_path, \"fasta\"):\n",
    "            sequences[record.id] = str(record.seq)\n",
    "        return sequences\n",
    "    \n",
    "    def load_seq(self, pdb_id):\n",
    "        return list(self.read_fasta_biopython(self.seq_dir + pdb_id + \".fasta\").values())[0]\n",
    "    \n",
    "    def first_load(self, idx):\n",
    "        pdb_id = self.get_pdb_id(idx)\n",
    "        feature = self.load_feature(pdb_id)\n",
    "        if self.is_train:\n",
    "            feature[\"seq\"] = self.load_seq(pdb_id)\n",
    "        else:\n",
    "            feature[\"seq\"] = None\n",
    "        return feature\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if idx in self.cache:\n",
    "            return self.cache[idx]\n",
    "        else:\n",
    "            data = self.first_load(idx)\n",
    "            self.cache[idx] = data\n",
    "            return data\n",
    "    \n",
    "    @staticmethod\n",
    "    def data_adapter(data):\n",
    "        coords = data[\"coords\"]\n",
    "        coords = np.concatenate((coords[\"P\"].reshape(-1, 1, 3), coords[\"C4'\"].reshape(-1, 1, 3), coords[\"N\"].reshape(-1, 1, 3)), axis=1)\n",
    "        seq = data[\"seq\"]\n",
    "        name = data[\"name\"]\n",
    "        return coords, seq, name\n",
    "\n",
    "\n",
    "    \n",
    "    def __iter__(self):\n",
    "        for idx in range(len(self)):\n",
    "            yield self.__getitem__(idx)\n",
    "\n",
    "    def get_lengths_by_indices(self, indices):\n",
    "        lengths = []\n",
    "        for idx in indices:\n",
    "            pdb_id = self.get_pdb_id(idx)\n",
    "            file_path = self.merge_coords_file_path(pdb_id)\n",
    "            with open(file_path, 'rb') as f:\n",
    "                # 读取文件头的前8字节（魔数和版本号）\n",
    "                version = np.lib.format.read_magic(f)\n",
    "                # 读取文件头信息（包含shape/dtype等）\n",
    "                shape, _, _ = np.lib.format._read_array_header(f, version)\n",
    "            lengths.append(shape[0])\n",
    "        return lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = RNADatasetV2(\n",
    "    data_path=\"/data/slz/sais_medicine/saisdata\",\n",
    "    is_train=True\n",
    ")\n",
    "# print(dataset[0]['coords'])\n",
    "coords, seq, name = dataset.data_adapter(dataset[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(76, 3, 3)"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coords.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct_data = construct_data_single(coords, seq)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'seq': tensor([1, 1, 1, 2, 3, 3, 1, 3, 0, 1, 2, 3, 2, 0, 1, 1, 3, 1, 1, 3, 3, 0, 1, 0,\n",
       "         1, 2, 1, 2, 0, 2, 2, 1, 2, 1, 0, 0, 0, 1, 2, 1, 1, 3, 1, 0, 1, 1, 3, 2,\n",
       "         1, 1, 3, 1, 1, 3, 3, 2, 0, 0, 1, 3, 2, 2, 0, 2, 3, 2, 0, 1, 1, 2, 2, 3,\n",
       "         0, 2, 2, 0]),\n",
       " 'coords': tensor([[[ 31.6120, -12.4040,   7.8380],\n",
       "          [ 34.1750, -12.5540,   4.9320],\n",
       "          [ 31.7160, -12.7520,   2.6360]],\n",
       " \n",
       "         [[ 35.0810,  -8.8840,   4.2240],\n",
       "          [ 35.1310,  -8.5690,   0.2950],\n",
       "          [ 31.8010,  -8.5760,  -0.1700]],\n",
       " \n",
       "         [[ 35.2750,  -4.6220,   0.3970],\n",
       "          [ 32.5750,  -3.3370,  -2.0970],\n",
       "          [ 29.5680,  -4.5040,  -1.3000]],\n",
       " \n",
       "         [[ 31.7710,   0.3900,  -2.4080],\n",
       "          [ 28.4640,   0.7440,  -4.5020],\n",
       "          [ 26.4250,  -0.5680,  -2.1280]],\n",
       " \n",
       "         [[ 27.9650,   4.4580,  -3.4810],\n",
       "          [ 24.0690,   4.5480,  -3.0020],\n",
       "          [ 24.1100,   2.6940,  -0.1420]],\n",
       " \n",
       "         [[ 23.6470,   8.0120,  -1.2640],\n",
       "          [ 20.3430,   7.5840,   0.8010],\n",
       "          [ 21.7470,   5.4740,   3.0720]],\n",
       " \n",
       "         [[ 20.4440,  10.7860,   3.0640],\n",
       "          [ 19.0130,   9.6100,   6.5420],\n",
       "          [ 21.9310,   7.6940,   7.3170]],\n",
       " \n",
       "         [[ 17.5500,  11.5430,   8.9260],\n",
       "          [ 19.1770,  14.9820,   9.7590],\n",
       "          [ 18.6100,  16.5790,   6.9870]],\n",
       " \n",
       "         [[ 15.4930,  15.8240,  10.9890],\n",
       "          [ 14.2450,  18.4280,  13.3390],\n",
       "          [ 12.7070,  20.5520,  10.9110]],\n",
       " \n",
       "         [[ 12.0060,  16.8750,  15.5710],\n",
       "          [  9.0420,  14.3320,  15.5720],\n",
       "          [  7.4270,  16.7810,  13.9730]],\n",
       " \n",
       "         [[ 10.1360,  11.9400,  12.9090],\n",
       "          [  7.3700,  11.2460,  10.2150],\n",
       "          [  7.3940,  14.4970,   9.1760]],\n",
       " \n",
       "         [[  9.6420,   9.7500,   7.4520],\n",
       "          [  8.2870,  11.0890,   3.9960],\n",
       "          [  9.0950,  14.3060,   4.6430]],\n",
       " \n",
       "         [[ 11.5570,  10.4110,   1.8890],\n",
       "          [ 12.6980,  13.9270,   0.7320],\n",
       "          [ 12.8970,  16.2460,   3.2000]],\n",
       " \n",
       "         [[ 16.0030,  15.2050,  -1.0200],\n",
       "          [ 15.6900,  19.0560,  -1.6970],\n",
       "          [ 16.4940,  20.5330,   1.2980]],\n",
       " \n",
       "         [[ 19.3170,  18.6970,  -2.9930],\n",
       "          [ 21.0870,  21.9210,  -1.5500],\n",
       "          [ 20.7990,  20.8840,   1.5980]],\n",
       " \n",
       "         [[ 24.7900,  20.7040,  -1.8610],\n",
       "          [ 25.3820,  20.2070,  -5.7300],\n",
       "          [ 25.8760,  23.4540,  -6.9500]],\n",
       " \n",
       "         [[ 28.9570,  18.8380,  -5.3330],\n",
       "          [ 31.6670,  21.5420,  -5.7250],\n",
       "          [ 31.0660,  23.8010,  -8.3410]],\n",
       " \n",
       "         [[ 34.0750,  21.9100,  -3.3870],\n",
       "          [ 32.0780,  22.2920,  -0.1170],\n",
       "          [ 35.2120,  23.0220,   1.3890]],\n",
       " \n",
       "         [[ 29.8940,  25.3310,   0.5080],\n",
       "          [ 30.7600,  28.1980,   3.0930],\n",
       "          [ 33.6580,  29.4590,   1.5200]],\n",
       " \n",
       "         [[ 28.3240,  30.4810,   3.8580],\n",
       "          [ 27.1390,  32.9340,   1.0460],\n",
       "          [ 26.7820,  30.4280,  -1.2100]],\n",
       " \n",
       "         [[ 23.3130,  33.1790,   1.9090],\n",
       "          [ 22.6530,  29.9100,   3.9420],\n",
       "          [ 23.3590,  28.2800,   0.9190]],\n",
       " \n",
       "         [[ 23.9830,  28.9120,   6.8790],\n",
       "          [ 21.1380,  26.8440,   7.1360],\n",
       "          [ 20.7880,  23.8790,   8.6380]],\n",
       " \n",
       "         [[ 18.2560,  28.9130,   7.5220],\n",
       "          [ 15.3590,  27.7870,   5.5410],\n",
       "          [ 15.5170,  24.4380,   5.4620]],\n",
       " \n",
       "         [[ 11.4070,  27.8010,   5.5860],\n",
       "          [ 10.2160,  26.2230,   2.1970],\n",
       "          [ 10.6400,  23.0740,   3.4080]],\n",
       " \n",
       "         [[  6.4480,  26.7250,   2.8700],\n",
       "          [  4.7300,  23.4930,   1.4700],\n",
       "          [  6.3450,  21.4570,   3.5990]],\n",
       " \n",
       "         [[  1.5800,  23.9480,   3.6560],\n",
       "          [  0.7680,  20.2300,   4.5530],\n",
       "          [  3.4500,  19.9040,   6.5820]],\n",
       " \n",
       "         [[ -1.8060,  20.9280,   7.3560],\n",
       "          [ -0.4590,  18.6000,  10.2120],\n",
       "          [  2.4100,  20.3350,  10.6590]],\n",
       " \n",
       "         [[ -2.4290,  20.1370,  13.1390],\n",
       "          [  0.0590,  20.4020,  16.1760],\n",
       "          [  1.8850,  22.9170,  14.7910]],\n",
       " \n",
       "         [[ -2.4720,  22.4260,  18.2220],\n",
       "          [ -0.5650,  25.1870,  20.2520],\n",
       "          [  0.6560,  26.9400,  17.5690]],\n",
       " \n",
       "         [[ -3.8810,  27.0170,  20.8000],\n",
       "          [ -2.8530,  30.7720,  20.9590],\n",
       "          [ -1.8980,  30.8870,  17.6720]],\n",
       " \n",
       "         [[ -6.3180,  32.5150,  20.4320],\n",
       "          [ -5.4180,  35.8080,  18.4800],\n",
       "          [ -4.6950,  34.3890,  15.4790]],\n",
       " \n",
       "         [[ -9.1660,  36.4990,  17.6400],\n",
       "          [ -9.5280,  37.5610,  13.9280],\n",
       "          [ -7.9330,  35.1440,  12.2010]],\n",
       " \n",
       "         [[-12.9660,  37.2150,  12.5950],\n",
       "          [-12.5040,  36.9520,   8.7390],\n",
       "          [-11.1930,  33.8090,   8.6180]],\n",
       " \n",
       "         [[-16.1680,  35.8490,   8.4090],\n",
       "          [-16.1010,  33.1420,   5.5600],\n",
       "          [-14.4720,  30.5910,   7.1210]],\n",
       " \n",
       "         [[-19.6510,  32.1570,   6.5020],\n",
       "          [-20.9190,  28.6570,   7.4080],\n",
       "          [-20.6430,  28.0550,   4.0670]],\n",
       " \n",
       "         [[-17.6750,  27.1200,   8.7540],\n",
       "          [-17.5050,  23.6200,   7.0160],\n",
       "          [-17.1280,  25.2770,   4.1150]],\n",
       " \n",
       "         [[-13.7920,  22.4910,   7.5320],\n",
       "          [-12.9020,  20.4220,   4.3150],\n",
       "          [-13.1530,  23.0960,   2.2720]],\n",
       " \n",
       "         [[ -9.1100,  20.5570,   4.9520],\n",
       "          [ -8.4140,  23.6940,   2.7420],\n",
       "          [ -9.9330,  25.7720,   4.8050]],\n",
       " \n",
       "         [[ -5.0210,  25.6000,   2.7700],\n",
       "          [ -5.9690,  29.2810,   1.7960],\n",
       "          [ -7.0980,  29.9810,   4.8770]],\n",
       " \n",
       "         [[ -2.4450,  30.7970,   2.4460],\n",
       "          [ -3.0230,  34.5540,   3.4900],\n",
       "          [ -4.1500,  33.7340,   6.5400]],\n",
       " \n",
       "         [[  0.6120,  35.2360,   4.5970],\n",
       "          [  0.2480,  37.4820,   7.8040],\n",
       "          [ -0.8240,  34.9990,   9.8420]],\n",
       " \n",
       "         [[  3.9860,  37.2700,   8.8850],\n",
       "          [  3.7120,  37.4810,  12.8010],\n",
       "          [  2.5030,  34.3220,  13.1210]],\n",
       " \n",
       "         [[  7.3980,  36.5910,  13.7780],\n",
       "          [  7.3060,  33.6330,  16.2150],\n",
       "          [  5.4250,  31.1170,  14.8970]],\n",
       " \n",
       "         [[ 10.4030,  31.7750,  17.3950],\n",
       "          [  9.0640,  28.7750,  19.4910],\n",
       "          [  7.8060,  27.1620,  16.7540]],\n",
       " \n",
       "         [[ 12.4440,  26.9810,  19.7510],\n",
       "          [ 11.3420,  23.3450,  18.8140],\n",
       "          [ 10.7330,  24.1410,  15.5490]],\n",
       " \n",
       "         [[ 14.7030,  21.9390,  19.0130],\n",
       "          [ 16.1520,  19.8790,  16.0790],\n",
       "          [ 17.8150,  21.9480,  13.7050]],\n",
       " \n",
       "         [[ 18.5990,  17.8740,  17.9590],\n",
       "          [ 22.2250,  17.7530,  18.7340],\n",
       "          [ 22.9580,  14.6450,  20.1010]],\n",
       " \n",
       "         [[ 23.4480,  18.0540,  15.2790],\n",
       "          [ 22.9000,  16.0760,  12.0770],\n",
       "          [ 23.1700,  18.9340,  10.0110]],\n",
       " \n",
       "         [[ 23.3880,  13.0010,  10.0060],\n",
       "          [ 22.4550,  11.1400,  13.3150],\n",
       "          [ 24.3060,   8.4620,  12.4870]],\n",
       " \n",
       "         [[ 25.2970,  12.9370,  15.2860],\n",
       "          [ 27.1300,  10.4240,  17.6930],\n",
       "          [ 28.6310,   8.8060,  15.1500]],\n",
       " \n",
       "         [[ 30.1420,  12.7090,  18.6080],\n",
       "          [ 33.1940,  10.2610,  18.9740],\n",
       "          [ 33.5410,   9.7560,  15.6650]],\n",
       " \n",
       "         [[ 35.8690,  13.1310,  19.0550],\n",
       "          [ 38.9470,  11.4410,  17.3140],\n",
       "          [ 37.6340,  11.5680,  14.2090]],\n",
       " \n",
       "         [[ 41.2180,  14.5900,  16.6950],\n",
       "          [ 42.5590,  14.3130,  13.0690],\n",
       "          [ 39.8160,  14.4150,  11.1210]],\n",
       " \n",
       "         [[ 44.5370,  17.1980,  11.7110],\n",
       "          [ 44.1810,  17.0500,   7.8210],\n",
       "          [ 40.9030,  17.9410,   7.7870]],\n",
       " \n",
       "         [[ 45.8560,  20.4370,   7.8180],\n",
       "          [ 43.7070,  22.6000,   5.4040],\n",
       "          [ 40.9920,  22.4510,   7.4240]],\n",
       " \n",
       "         [[ 45.1490,  25.8970,   6.7970],\n",
       "          [ 43.0650,  28.6410,   8.6420],\n",
       "          [ 41.4370,  29.4530,   5.7940]],\n",
       " \n",
       "         [[ 40.6120,  26.0460,  10.0470],\n",
       "          [ 37.0660,  27.4340,   9.0830],\n",
       "          [ 37.5340,  27.0040,   5.7620]],\n",
       " \n",
       "         [[ 35.2430,  24.1200,   9.8660],\n",
       "          [ 32.1610,  22.9060,   7.9600],\n",
       "          [ 34.2730,  20.4730,   6.4700]],\n",
       " \n",
       "         [[ 29.2490,  21.9190,   9.3440],\n",
       "          [ 27.2580,  18.6500,   8.5410],\n",
       "          [ 25.8870,  20.0950,   5.8340]],\n",
       " \n",
       "         [[ 30.3680,  16.9430,   6.9760],\n",
       "          [ 29.3480,  16.7160,   3.1760],\n",
       "          [ 29.3390,  20.2670,   3.0300]],\n",
       " \n",
       "         [[ 30.9610,  14.9690,   0.6330],\n",
       "          [ 34.8410,  15.3140,   0.1430],\n",
       "          [ 35.6650,  16.2040,   3.3230]],\n",
       " \n",
       "         [[ 35.1240,  11.4540,   0.5710],\n",
       "          [ 38.6650,  10.7880,   2.1400],\n",
       "          [ 37.8670,  11.8990,   5.2590]],\n",
       " \n",
       "         [[ 37.9330,   7.0250,   2.6660],\n",
       "          [ 39.6850,   6.0340,   6.0440],\n",
       "          [ 37.7450,   7.9490,   8.0400]],\n",
       " \n",
       "         [[ 37.7450,   2.6900,   6.8120],\n",
       "          [ 37.8810,   2.5400,  10.7380],\n",
       "          [ 35.6930,   5.0720,  11.1270]],\n",
       " \n",
       "         [[ 35.2420,  -0.2770,  10.9980],\n",
       "          [ 33.1870,   0.9910,  14.0740],\n",
       "          [ 31.6840,   3.4660,  12.3200]],\n",
       " \n",
       "         [[ 30.1990,  -1.5250,  13.9370],\n",
       "          [ 27.1530,   0.3870,  15.5590],\n",
       "          [ 26.5760,   2.4510,  12.9590]],\n",
       " \n",
       "         [[ 24.5270,  -2.3250,  14.4260],\n",
       "          [ 21.2620,  -0.3140,  13.4890],\n",
       "          [ 22.6000,   1.3240,  10.8110]],\n",
       " \n",
       "         [[ 19.3100,  -2.9660,  11.3970],\n",
       "          [ 16.9350,  -0.9690,   8.9700],\n",
       "          [ 19.3070,   0.3690,   6.9500]],\n",
       " \n",
       "         [[ 16.2190,  -3.9540,   6.4900],\n",
       "          [ 15.6230,  -2.1680,   3.0130],\n",
       "          [ 18.8770,  -1.4370,   2.6090]],\n",
       " \n",
       "         [[ 15.5970,  -5.3430,   0.8480],\n",
       "          [ 17.5820,  -4.3320,  -2.3810],\n",
       "          [ 20.6500,  -4.0420,  -0.9070]],\n",
       " \n",
       "         [[ 18.4600,  -7.7200,  -4.1570],\n",
       "          [ 21.6030,  -6.9590,  -6.3950],\n",
       "          [ 23.8250,  -6.8570,  -3.8380]],\n",
       " \n",
       "         [[ 22.2000, -10.7010,  -7.2560],\n",
       "          [ 26.1190, -10.9790,  -7.1530],\n",
       "          [ 26.2980, -10.4810,  -3.8310]],\n",
       " \n",
       "         [[ 26.9420, -14.7960,  -7.0250],\n",
       "          [ 30.5320, -15.0990,  -5.4620],\n",
       "          [ 29.7820, -14.7130,  -2.1890]],\n",
       " \n",
       "         [[ 29.9920, -18.9560,  -5.5270],\n",
       "          [ 31.7710, -20.1170,  -2.2250],\n",
       "          [ 29.3890, -18.8520,  -0.1920]],\n",
       " \n",
       "         [[ 30.5810, -23.6620,  -2.6930],\n",
       "          [ 33.5880, -24.0510,  -5.0910],\n",
       "          [ 35.6700, -23.3390,  -2.3290]],\n",
       " \n",
       "         [[ 34.8470, -27.4740,  -6.2560],\n",
       "          [ 34.7950, -28.1210,  -9.7740],\n",
       "          [ 32.7870, -28.3430, -12.7690]]]),\n",
       " 'node_s': tensor([[0., 1., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 1., 0., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [0., 0., 0., 1.],\n",
       "         [1., 0., 0., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [0., 0., 1., 0.],\n",
       "         [1., 0., 0., 0.]]),\n",
       " 'node_v': tensor([[[ 1.5448e-01,  6.4395e-01, -7.4931e-01],\n",
       "          [ 0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
       "          [-7.2965e-01, -5.8752e-02, -6.8129e-01],\n",
       "          [-6.6096e-01,  3.8683e-02,  7.4942e-01]],\n",
       " \n",
       "         [[-4.0603e-01,  8.3112e-01, -3.7998e-01],\n",
       "          [-1.5448e-01, -6.4395e-01,  7.4931e-01],\n",
       "          [-9.9039e-01, -2.0819e-03, -1.3830e-01],\n",
       "          [-1.2684e-02, -7.9910e-02,  9.9672e-01]],\n",
       " \n",
       "         [[-6.5545e-01,  6.5066e-01, -3.8345e-01],\n",
       "          [ 4.0603e-01, -8.3112e-01,  3.7998e-01],\n",
       "          [-9.0504e-01, -3.5124e-01,  2.3988e-01],\n",
       "          [ 6.9342e-01, -3.3002e-01,  6.4051e-01]],\n",
       " \n",
       "         [[-7.3213e-01,  6.3368e-01,  2.4987e-01],\n",
       "          [ 6.5545e-01, -6.5066e-01,  3.8345e-01],\n",
       "          [-6.0088e-01, -3.8664e-01,  6.9961e-01],\n",
       "          [ 8.4143e-01, -9.0072e-02,  5.3280e-01]],\n",
       " \n",
       "         [[-6.0794e-01,  4.9536e-01,  6.2050e-01],\n",
       "          [ 7.3213e-01, -6.3368e-01, -2.4987e-01],\n",
       "          [ 1.2028e-02, -5.4392e-01,  8.3905e-01],\n",
       "          [ 9.9227e-01, -2.2922e-02, -1.2200e-01]],\n",
       " \n",
       "         [[-2.1343e-01,  3.2512e-01,  9.2127e-01],\n",
       "          [ 6.0794e-01, -4.9536e-01, -6.2050e-01],\n",
       "          [ 4.1257e-01, -6.2003e-01,  6.6734e-01],\n",
       "          [ 8.4293e-01,  1.0919e-01, -5.2683e-01]],\n",
       " \n",
       "         [[ 2.6182e-02,  8.5764e-01,  5.1359e-01],\n",
       "          [ 2.1343e-01, -3.2512e-01, -9.2127e-01],\n",
       "          [ 8.1604e-01, -5.3582e-01,  2.1673e-01],\n",
       "          [ 3.6316e-01,  2.9844e-01, -8.8264e-01]],\n",
       " \n",
       "         [[-7.0446e-01,  4.9221e-01,  5.1135e-01],\n",
       "          [-2.6182e-02, -8.5764e-01, -5.1359e-01],\n",
       "          [-1.7452e-01,  4.9154e-01, -8.5319e-01],\n",
       "          [-4.1776e-01, -8.8302e-01, -2.1389e-01]],\n",
       " \n",
       "         [[-7.4454e-01, -5.8613e-01,  3.1954e-01],\n",
       "          [ 7.0446e-01, -4.9221e-01, -5.1135e-01],\n",
       "          [-4.3035e-01,  5.9433e-01, -6.7939e-01],\n",
       "          [ 3.3521e-01, -6.9943e-01, -6.3121e-01]],\n",
       " \n",
       "         [[-2.6107e-01, -4.8186e-01, -8.3646e-01],\n",
       "          [ 7.4454e-01,  5.8613e-01, -3.1954e-01],\n",
       "          [-4.8338e-01,  7.3300e-01, -4.7859e-01],\n",
       "          [ 7.5895e-01,  6.5115e-01, -2.5592e-04]],\n",
       " \n",
       "         [[ 1.4583e-01, -2.4967e-02, -9.8899e-01],\n",
       "          [ 2.6107e-01,  4.8186e-01,  8.3646e-01],\n",
       "          [ 7.0318e-03,  9.5251e-01, -3.0442e-01],\n",
       "          [ 7.0507e-01,  1.7690e-01,  6.8672e-01]],\n",
       " \n",
       "         [[ 7.1401e-01,  4.5939e-01, -5.2835e-01],\n",
       "          [-1.4583e-01,  2.4967e-02,  9.8899e-01],\n",
       "          [ 2.3909e-01,  9.5194e-01,  1.9145e-01],\n",
       "          [ 3.4336e-01, -3.3931e-01,  8.7577e-01]],\n",
       " \n",
       "         [[ 4.6637e-01,  7.9947e-01, -3.7861e-01],\n",
       "          [-7.1401e-01, -4.5939e-01,  5.2835e-01],\n",
       "          [ 5.8661e-02,  6.8359e-01,  7.2751e-01],\n",
       "          [-2.9458e-01, -9.0774e-01,  2.9871e-01]],\n",
       " \n",
       "         [[ 8.8301e-01,  4.6874e-01,  2.4051e-02],\n",
       "          [-4.6637e-01, -7.9947e-01,  3.7861e-01],\n",
       "          [ 2.3407e-01,  4.3001e-01,  8.7195e-01],\n",
       "          [ 7.9795e-02, -9.8176e-01,  1.7259e-01]],\n",
       " \n",
       "         [[ 6.8901e-01, -2.7496e-01, -6.7056e-01],\n",
       "          [-8.8301e-01, -4.6874e-01, -2.4051e-02],\n",
       "          [-8.6567e-02, -3.1170e-01,  9.4623e-01],\n",
       "          [-4.4800e-01, -8.1602e-01, -3.6524e-01]],\n",
       " \n",
       "         [[ 9.7818e-01,  2.0777e-01,  7.7820e-04],\n",
       "          [-6.8901e-01,  2.7496e-01,  6.7056e-01],\n",
       "          [ 1.4100e-01,  9.2675e-01, -3.4821e-01],\n",
       "          [-1.5005e-01,  1.2597e-01,  9.8062e-01]],\n",
       " \n",
       "         [[ 7.2450e-02,  1.3221e-01,  9.8857e-01],\n",
       "          [-9.7818e-01, -2.0777e-01, -7.7820e-04],\n",
       "          [-1.7131e-01,  6.4391e-01, -7.4567e-01],\n",
       "          [-7.0421e-01, -7.0265e-01,  1.0186e-01]],\n",
       " \n",
       "         [[-1.9241e-01,  8.6219e-01,  4.6862e-01],\n",
       "          [-7.2450e-02, -1.3221e-01, -9.8857e-01],\n",
       "          [ 8.8210e-01,  2.0547e-01,  4.2388e-01],\n",
       "          [ 5.1863e-01, -9.9206e-02, -8.4923e-01]],\n",
       " \n",
       "         [[-5.7446e-01,  7.5135e-01, -3.2475e-01],\n",
       "          [ 1.9241e-01, -8.6219e-01, -4.6862e-01],\n",
       "          [ 8.2090e-01,  3.5720e-01, -4.4557e-01],\n",
       "          [-2.1889e-01, -7.2468e-01, -6.5340e-01]],\n",
       " \n",
       "         [[-7.3105e-01, -4.9280e-01,  4.7194e-01],\n",
       "          [ 5.7446e-01, -7.5135e-01,  3.2475e-01],\n",
       "          [-1.0529e-01, -7.3907e-01, -6.6534e-01],\n",
       "          [ 3.0267e-01, -6.2653e-01,  7.1823e-01]],\n",
       " \n",
       "         [[-3.2376e-01, -6.5521e-01,  6.8256e-01],\n",
       "          [ 7.3105e-01,  4.9280e-01, -4.7194e-01],\n",
       "          [ 2.0135e-01, -4.6488e-01, -8.6217e-01],\n",
       "          [ 1.6898e-01,  8.3697e-01, -5.2051e-01]],\n",
       " \n",
       "         [[-9.5225e-01,  1.5539e-01, -2.6282e-01],\n",
       "          [ 3.2376e-01,  6.5521e-01, -6.8256e-01],\n",
       "          [-1.0472e-01, -8.8716e-01,  4.4942e-01],\n",
       "          [ 8.0673e-01,  5.8641e-01, -7.2875e-02]],\n",
       " \n",
       "         [[-8.1238e-01, -2.4705e-01, -5.2821e-01],\n",
       "          [ 9.5225e-01, -1.5539e-01,  2.6282e-01],\n",
       "          [ 4.7113e-02, -9.9861e-01, -2.3556e-02],\n",
       "          [ 7.8600e-01,  3.0550e-01,  5.3747e-01]],\n",
       " \n",
       "         [[-8.8904e-01, -4.4241e-01, -1.1781e-01],\n",
       "          [ 8.1238e-01,  2.4705e-01,  5.2821e-01],\n",
       "          [ 1.2469e-01, -9.2608e-01,  3.5614e-01],\n",
       "          [ 3.0356e-01,  4.0219e-01,  8.6377e-01]],\n",
       " \n",
       "         [[-6.6172e-01, -5.4497e-01,  5.1491e-01],\n",
       "          [ 8.8904e-01,  4.4241e-01,  1.1781e-01],\n",
       "          [ 4.8073e-01, -6.0604e-01,  6.3373e-01],\n",
       "          [ 4.3839e-01,  8.2473e-01,  3.5725e-01]],\n",
       " \n",
       "         [[-2.0397e-01, -2.7096e-01,  9.4073e-01],\n",
       "          [ 6.6172e-01,  5.4497e-01, -5.1491e-01],\n",
       "          [ 7.9377e-01, -9.6484e-02,  6.0051e-01],\n",
       "          [ 2.0768e-01,  9.5091e-01, -2.2942e-01]],\n",
       " \n",
       "         [[ 8.2856e-02,  2.8824e-01,  9.5397e-01],\n",
       "          [ 2.0397e-01,  2.7096e-01, -9.4073e-01],\n",
       "          [ 8.4819e-01,  5.1294e-01,  1.3215e-01],\n",
       "          [-3.4335e-01,  5.9341e-01, -7.2800e-01]],\n",
       " \n",
       "         [[-9.8787e-02,  7.5753e-01,  6.4528e-01],\n",
       "          [-8.2856e-02, -2.8824e-01, -9.5397e-01],\n",
       "          [ 5.3665e-01,  7.3914e-01, -4.0704e-01],\n",
       "          [-6.3228e-01, -6.7346e-02, -7.7180e-01]],\n",
       " \n",
       "         [[-3.7652e-01,  9.1908e-01,  1.1634e-01],\n",
       "          [ 9.8787e-02, -7.5753e-01, -6.4528e-01],\n",
       "          [ 3.5602e-01,  5.1113e-01, -7.8230e-01],\n",
       "          [-4.8625e-01, -7.0401e-01, -5.1762e-01]],\n",
       " \n",
       "         [[-4.1563e-01,  8.1602e-01, -4.0169e-01],\n",
       "          [ 3.7652e-01, -9.1908e-01, -1.1634e-01],\n",
       "          [ 2.7884e-01,  3.3578e-02, -9.5975e-01],\n",
       "          [-2.6383e-01, -9.6371e-01, -4.0807e-02]],\n",
       " \n",
       "         [[-6.4435e-01,  2.7483e-01, -7.1364e-01],\n",
       "          [ 4.1563e-01, -8.1602e-01,  4.0169e-01],\n",
       "          [ 2.1281e-01, -4.1767e-01, -8.8332e-01],\n",
       "          [-2.2886e-01, -8.3739e-01,  4.9638e-01]],\n",
       " \n",
       "         [[-4.9495e-01, -1.0128e-01, -8.6300e-01],\n",
       "          [ 6.4435e-01, -2.7483e-01,  7.1364e-01],\n",
       "          [ 4.7305e-01, -7.1685e-01, -5.1220e-01],\n",
       "          [ 9.3350e-02, -2.7386e-01,  9.5723e-01]],\n",
       " \n",
       "         [[-5.8691e-01, -6.2167e-01, -5.1871e-01],\n",
       "          [ 4.9495e-01,  1.0128e-01,  8.6300e-01],\n",
       "          [ 3.8473e-01, -9.2235e-01, -3.5509e-02],\n",
       "          [-1.1869e-01,  6.7566e-02,  9.9063e-01]],\n",
       " \n",
       "         [[-7.0470e-01, -6.5600e-01,  2.7030e-01],\n",
       "          [ 5.8691e-01,  6.2167e-01,  5.1871e-01],\n",
       "          [ 4.7833e-01, -7.4907e-01,  4.5837e-01],\n",
       "          [-1.7046e-02,  6.8871e-01,  7.2484e-01]],\n",
       " \n",
       "         [[ 5.5989e-01, -8.2607e-01, -6.4288e-02],\n",
       "          [ 7.0470e-01,  6.5600e-01, -2.7030e-01],\n",
       "          [ 8.1034e-02, -1.7675e-01, -9.8092e-01],\n",
       "          [ 3.3096e-01,  9.1353e-01, -2.3647e-01]],\n",
       " \n",
       "         [[ 7.3982e-01, -5.1400e-01, -4.3412e-01],\n",
       "          [-5.5989e-01,  8.2607e-01,  6.4288e-02],\n",
       "          [ 1.1213e-01,  4.9285e-01, -8.6286e-01],\n",
       "          [-4.3462e-02,  8.9481e-01,  4.4433e-01]],\n",
       " \n",
       "         [[ 7.7747e-01,  5.6682e-01, -2.7250e-01],\n",
       "          [-7.3982e-01,  5.1400e-01,  4.3412e-01],\n",
       "          [-7.4382e-02,  7.9242e-01, -6.0543e-01],\n",
       "          [-2.2663e-01,  5.2685e-01,  8.1918e-01]],\n",
       " \n",
       "         [[ 3.9618e-01,  9.0529e-01, -1.5329e-01],\n",
       "          [-7.7747e-01, -5.6682e-01,  2.7250e-01],\n",
       "          [-4.6048e-01,  6.2995e-01,  6.2540e-01],\n",
       "          [-1.7847e-01, -8.0438e-01,  5.6668e-01]],\n",
       " \n",
       "         [[ 4.6962e-01,  8.4056e-01,  2.7004e-01],\n",
       "          [-3.9618e-01, -9.0529e-01,  1.5329e-01],\n",
       "          [-3.3650e-01,  2.0863e-01,  9.1828e-01],\n",
       "          [ 2.4160e-01, -9.3809e-01,  2.4822e-01]],\n",
       " \n",
       "         [[ 5.3144e-01,  4.7572e-01,  7.0090e-01],\n",
       "          [-4.6962e-01, -8.4056e-01, -2.7004e-01],\n",
       "          [-3.3608e-01, -2.4453e-01,  9.0954e-01],\n",
       "          [ 1.4663e-01, -9.5308e-01, -2.6484e-01]],\n",
       " \n",
       "         [[ 5.6971e-01, -1.6438e-04,  8.2184e-01],\n",
       "          [-5.3144e-01, -4.7572e-01, -7.0090e-01],\n",
       "          [-3.1656e-01, -7.3322e-01,  6.0181e-01],\n",
       "          [ 9.2570e-02, -5.7119e-01, -8.1558e-01]],\n",
       " \n",
       "         [[ 5.7272e-01, -6.1320e-01,  5.4404e-01],\n",
       "          [-5.6971e-01,  1.6438e-04, -8.2184e-01],\n",
       "          [-3.5584e-01, -9.2979e-01,  9.4186e-02],\n",
       "          [ 6.9698e-02, -5.3672e-02, -9.9612e-01]],\n",
       " \n",
       "         [[ 2.8738e-01, -7.9412e-01,  5.3552e-01],\n",
       "          [-5.7272e-01,  6.1320e-01, -5.4404e-01],\n",
       "          [-5.5215e-01, -7.3855e-01, -3.8689e-01],\n",
       "          [ 2.3998e-02,  7.7158e-01, -6.3568e-01]],\n",
       " \n",
       "         [[ 3.8433e-01, -9.1610e-01, -1.1422e-01],\n",
       "          [-2.8738e-01,  7.9412e-01, -5.3552e-01],\n",
       "          [-3.6817e-01, -4.7206e-01, -8.0101e-01],\n",
       "          [ 3.4360e-01,  7.6984e-01, -5.3786e-01]],\n",
       " \n",
       "         [[ 7.3670e-01, -5.3085e-01, -4.1889e-01],\n",
       "          [-3.8433e-01,  9.1610e-01,  1.1422e-01],\n",
       "          [-1.7831e-01,  2.3306e-01, -9.5597e-01],\n",
       "          [ 2.8161e-01,  9.2917e-01,  2.3945e-01]],\n",
       " \n",
       "         [[ 8.7248e-01, -3.0543e-01,  3.8143e-01],\n",
       "          [-7.3670e-01,  5.3085e-01,  4.1889e-01],\n",
       "          [ 4.6698e-01,  5.8098e-01, -6.6663e-01],\n",
       "          [-3.7474e-01,  5.3275e-01,  7.5878e-01]],\n",
       " \n",
       "         [[ 9.7853e-02, -2.4311e-01, -9.6505e-01],\n",
       "          [-8.7248e-01,  3.0543e-01, -3.8143e-01],\n",
       "          [ 2.1102e-01, -8.9476e-01,  3.9354e-01],\n",
       "          [-9.7739e-01,  3.2616e-02, -2.0890e-01]],\n",
       " \n",
       "         [[-8.7113e-02, -9.6627e-01,  2.4235e-01],\n",
       "          [-9.7853e-02,  2.4311e-01,  9.6505e-01],\n",
       "          [ 7.6339e-02,  8.0806e-01, -5.8413e-01],\n",
       "          [ 1.4408e-01,  5.2007e-01,  8.4189e-01]],\n",
       " \n",
       "         [[ 7.2539e-01, -1.1110e-01,  6.7931e-01],\n",
       "          [ 8.7113e-02,  9.6627e-01, -2.4235e-01],\n",
       "          [ 5.5104e-01, -7.9724e-01, -2.4650e-01],\n",
       "          [ 2.3866e-01,  4.7603e-01, -8.4642e-01]],\n",
       " \n",
       "         [[ 9.7807e-01, -2.6290e-02,  2.0661e-01],\n",
       "          [-7.2539e-01,  1.1110e-01, -6.7931e-01],\n",
       "          [ 4.4578e-01, -4.8052e-01, -7.5524e-01],\n",
       "          [-4.6605e-01,  6.3895e-01, -6.1200e-01]],\n",
       " \n",
       "         [[ 9.4267e-01,  1.9335e-01, -2.7200e-01],\n",
       "          [-9.7807e-01,  2.6290e-02, -2.0661e-01],\n",
       "          [ 1.0311e-01, -1.5006e-01, -9.8328e-01],\n",
       "          [-7.7668e-01,  6.2297e-01, -9.3141e-02]],\n",
       " \n",
       "         [[ 5.7606e-01,  4.5804e-01, -6.7702e-01],\n",
       "          [-9.4267e-01, -1.9335e-01,  2.7200e-01],\n",
       "          [-3.8920e-01,  3.7645e-02, -9.2038e-01],\n",
       "          [-7.8534e-01,  4.3119e-01,  4.4421e-01]],\n",
       " \n",
       "         [[ 2.6430e-01,  4.4598e-01, -8.5513e-01],\n",
       "          [-5.7606e-01, -4.5804e-01,  6.7702e-01],\n",
       "          [-8.1494e-01,  3.0304e-02, -5.7875e-01],\n",
       "          [-3.4598e-01,  7.1467e-02,  9.3552e-01]],\n",
       " \n",
       "         [[-7.8063e-02,  9.1403e-01, -3.9806e-01],\n",
       "          [-2.6430e-01, -4.4598e-01,  8.5513e-01],\n",
       "          [-9.6494e-01,  2.6228e-01, -1.0008e-02],\n",
       "          [ 9.1070e-02,  3.7861e-02,  9.9512e-01]],\n",
       " \n",
       "         [[-9.3259e-02,  8.7753e-01,  4.7036e-01],\n",
       "          [ 7.8063e-02, -9.1403e-01,  3.9806e-01],\n",
       "          [-8.0152e-01, -4.3988e-02,  5.9634e-01],\n",
       "          [ 5.5259e-01, -5.5619e-01,  6.2073e-01]],\n",
       " \n",
       "         [[-9.7782e-01, -1.9674e-01,  7.1882e-02],\n",
       "          [ 9.3259e-02, -8.7753e-01, -4.7036e-01],\n",
       "          [-4.8173e-01,  2.4027e-01, -8.4274e-01],\n",
       "          [ 5.3319e-01, -7.0205e-01, -4.7204e-01]],\n",
       " \n",
       "         [[-7.2460e-01, -6.6891e-01, -1.6590e-01],\n",
       "          [ 9.7782e-01,  1.9674e-01, -7.1882e-02],\n",
       "          [ 1.3841e-01, -1.2717e-01, -9.8218e-01],\n",
       "          [ 9.0273e-01, -3.5335e-01,  2.4541e-01]],\n",
       " \n",
       "         [[-7.5217e-01, -6.5291e-01,  8.9131e-02],\n",
       "          [ 7.2460e-01,  6.6891e-01,  1.6590e-01],\n",
       "          [ 5.9499e-01, -6.8542e-01, -4.1976e-01],\n",
       "          [ 8.0645e-01,  3.1766e-01,  4.9873e-01]],\n",
       " \n",
       "         [[ 3.4410e-01, -3.1841e-01, -8.8329e-01],\n",
       "          [ 7.5217e-01,  6.5291e-01, -8.9131e-02],\n",
       "          [-4.0793e-01,  4.2995e-01, -8.0544e-01],\n",
       "          [ 5.0909e-01,  8.3587e-01,  2.0532e-01]],\n",
       " \n",
       "         [[ 8.5435e-01, -2.1806e-01, -4.7174e-01],\n",
       "          [-3.4410e-01,  3.1841e-01,  8.8329e-01],\n",
       "          [-2.5320e-03,  9.9915e-01, -4.1080e-02],\n",
       "          [ 2.5881e-01,  5.7599e-02,  9.6421e-01]],\n",
       " \n",
       "         [[ 6.1158e-01, -7.2385e-01,  3.1938e-01],\n",
       "          [-8.5435e-01,  2.1806e-01,  4.7174e-01],\n",
       "          [ 2.4211e-01,  2.6150e-01,  9.3435e-01],\n",
       "          [-9.8828e-01, -8.7876e-02,  1.2481e-01]],\n",
       " \n",
       "         [[ 1.6358e-01, -7.6240e-01,  6.2609e-01],\n",
       "          [-6.1158e-01,  7.2385e-01, -3.1938e-01],\n",
       "          [-2.3431e-01,  3.2621e-01,  9.1580e-01],\n",
       "          [-9.0104e-01,  1.6947e-01, -3.9925e-01]],\n",
       " \n",
       "         [[-2.9461e-01, -5.7060e-01,  7.6657e-01],\n",
       "          [-1.6358e-01,  7.6240e-01, -6.2609e-01],\n",
       "          [-5.7421e-01,  5.6681e-01,  5.9078e-01],\n",
       "          [-4.4555e-01,  2.5202e-01, -8.5905e-01]],\n",
       " \n",
       "         [[-7.8714e-01, -2.5975e-01,  5.5941e-01],\n",
       "          [ 2.9461e-01,  5.7060e-01, -7.6657e-01],\n",
       "          [-6.4946e-01,  7.5157e-01,  1.1547e-01],\n",
       "          [-3.4595e-02,  3.8156e-02, -9.9867e-01]],\n",
       " \n",
       "         [[-9.6647e-01, -9.6743e-02,  2.3785e-01],\n",
       "          [ 7.8714e-01,  2.5975e-01, -5.5941e-01],\n",
       "          [-4.4396e-01,  7.3107e-01, -5.1810e-01],\n",
       "          [ 5.2550e-01, -3.2425e-01, -7.8658e-01]],\n",
       " \n",
       "         [[-9.3756e-01, -1.1157e-01, -3.2944e-01],\n",
       "          [ 9.6647e-01,  9.6743e-02, -2.3785e-01],\n",
       "          [-1.7125e-01,  6.1257e-01, -7.7165e-01],\n",
       "          [ 7.7207e-01, -4.8464e-01, -4.1113e-01]],\n",
       " \n",
       "         [[-6.8784e-01, -1.0412e-01, -7.1836e-01],\n",
       "          [ 9.3756e-01,  1.1157e-01,  3.2944e-01],\n",
       "          [ 3.9209e-01,  4.8000e-01, -7.8477e-01],\n",
       "          [ 8.2712e-01, -5.0944e-01,  2.3737e-01]],\n",
       " \n",
       "         [[-2.1105e-01, -1.9287e-01, -9.5826e-01],\n",
       "          [ 6.8784e-01,  1.0412e-01,  7.1836e-01],\n",
       "          [ 6.9955e-01,  3.9461e-01, -5.9574e-01],\n",
       "          [ 6.0288e-01, -5.0693e-01,  6.1608e-01]],\n",
       " \n",
       "         [[ 3.1941e-01, -3.5284e-01, -8.7948e-01],\n",
       "          [ 2.1105e-01,  1.9287e-01,  9.5826e-01],\n",
       "          [ 9.6860e-01,  2.1759e-01, -1.2026e-01],\n",
       "          [ 1.5073e-01, -4.5169e-01,  8.7935e-01]],\n",
       " \n",
       "         [[ 6.4238e-01, -4.1968e-01, -6.4126e-01],\n",
       "          [-3.1941e-01,  3.5284e-01,  8.7948e-01],\n",
       "          [ 8.9811e-01,  8.4893e-02,  4.3149e-01],\n",
       "          [-5.0601e-01, -2.5772e-01,  8.2313e-01]],\n",
       " \n",
       "         [[ 7.4113e-01, -6.5973e-01, -1.2440e-01],\n",
       "          [-6.4238e-01,  4.1968e-01,  6.4126e-01],\n",
       "          [ 6.5563e-01,  3.0097e-02,  7.5448e-01],\n",
       "          [-7.9919e-01, -1.9351e-01,  5.6907e-01]],\n",
       " \n",
       "         [[ 7.0387e-01, -6.5713e-01,  2.6971e-01],\n",
       "          [-7.4113e-01,  6.5973e-01,  1.2440e-01],\n",
       "          [ 5.3213e-02,  1.4804e-01,  9.8755e-01],\n",
       "          [-9.9715e-01,  7.0734e-02, -2.6207e-02]],\n",
       " \n",
       "         [[ 2.0316e-01, -8.2280e-01,  5.3077e-01],\n",
       "          [-7.0387e-01,  6.5713e-01, -2.6971e-01],\n",
       "          [-2.2190e-01,  1.1420e-01,  9.6836e-01],\n",
       "          [-9.1414e-01,  7.7154e-02, -3.9799e-01]],\n",
       " \n",
       "         [[ 3.4974e-01, -7.5721e-01, -5.5165e-01],\n",
       "          [-2.0316e-01,  8.2280e-01, -5.3077e-01],\n",
       "          [-7.0526e-01,  3.7454e-01,  6.0193e-01],\n",
       "          [-4.5310e-01,  2.9570e-01, -8.4099e-01]],\n",
       " \n",
       "         [[ 1.9096e-01, -6.4391e-01, -7.4089e-01],\n",
       "          [-3.4974e-01,  7.5721e-01,  5.5165e-01],\n",
       "          [ 5.8958e-01,  2.0162e-01,  7.8214e-01],\n",
       "          [-7.7786e-01,  1.0063e-01,  6.2032e-01]],\n",
       " \n",
       "         [[ 0.0000e+00,  0.0000e+00,  0.0000e+00],\n",
       "          [-1.9096e-01,  6.4391e-01,  7.4089e-01],\n",
       "          [-5.5582e-01, -6.1450e-02, -8.2903e-01],\n",
       "          [ 1.4536e-02,  1.8086e-01,  9.8340e-01]]]),\n",
       " 'edge_s': tensor([[ 2.2683e-11,  2.8078e-07,  3.5709e-04,  ..., -3.1623e-03,\n",
       "          -1.0000e-03, -3.1623e-04],\n",
       "         [ 8.7764e-39,  1.3247e-30,  2.0544e-23,  ..., -6.3245e-03,\n",
       "          -2.0000e-03, -6.3246e-04],\n",
       "         [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ..., -2.0101e-01,\n",
       "          -6.3956e-02, -2.0237e-02],\n",
       "         ...,\n",
       "         [ 0.0000e+00,  0.0000e+00,  9.2404e-39,  ...,  9.4867e-03,\n",
       "           3.0000e-03,  9.4868e-04],\n",
       "         [ 6.4898e-37,  5.9621e-29,  5.6274e-22,  ...,  6.3245e-03,\n",
       "           2.0000e-03,  6.3246e-04],\n",
       "         [ 7.8585e-12,  1.2194e-07,  1.9440e-04,  ...,  3.1623e-03,\n",
       "           1.0000e-03,  3.1623e-04]]),\n",
       " 'edge_v': tensor([[[-0.1545, -0.6439,  0.7493]],\n",
       " \n",
       "         [[ 0.1367, -0.7877,  0.6007]],\n",
       " \n",
       "         [[ 0.0603, -0.8274, -0.5584]],\n",
       " \n",
       "         ...,\n",
       " \n",
       "         [[ 0.2968, -0.9065, -0.3002]],\n",
       " \n",
       "         [[ 0.2650, -0.7015, -0.6616]],\n",
       " \n",
       "         [[ 0.1910, -0.6439, -0.7409]]]),\n",
       " 'edge_index': tensor([[ 0,  0,  0,  ..., 75, 75, 75],\n",
       "         [ 1,  2, 64,  ..., 72, 73, 74]]),\n",
       " 'mask': tensor([True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True])}"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "struct_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNADatasetV3(RNADatasetV2):\n",
    "    def __init__(self, data_path, is_train=True):\n",
    "        super().__init__(data_path, is_train)\n",
    "        \n",
    "    def __getitem__(self, idx):\n",
    "        data = super().__getitem__(idx)\n",
    "        coords, seq, name = self.data_adapter(data)\n",
    "        struct_data = construct_data_single(coords, seq)\n",
    "        # print(struct_data)\n",
    "        zt = torch.randn(struct_data['node_s'].shape)\n",
    "        struct_data['z_t'] = zt\n",
    "        return struct_data\n",
    "    \n",
    "    @staticmethod\n",
    "    def data_adapter(data):\n",
    "        coords = data[\"coords\"]\n",
    "        coords = np.concatenate((coords[\"P\"].reshape(-1, 1, 3), coords[\"C4'\"].reshape(-1, 1, 3), coords[\"N\"].reshape(-1, 1, 3)), axis=1)\n",
    "        seq = data[\"seq\"]\n",
    "        name = data[\"name\"]\n",
    "        return coords, seq, name\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([76, 4])"
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = RNADatasetV3(\n",
    "    data_path=\"/data/slz/sais_medicine/saisdata\",\n",
    "    is_train=True\n",
    ")\n",
    "dataset[0]['z_t'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.data import Batch\n",
    "\n",
    "@torch.no_grad()\n",
    "def geo_batch(batch):\n",
    "    data_list = []\n",
    "    # print(len(batch['z_t']))\n",
    "    batch_size, length = batch['node_s'].shape[:2]\n",
    "\n",
    "    for i in range(batch_size):\n",
    "        data_list.append(torch_geometric.data.Data(\n",
    "            z_t=batch['z_t'][i],\n",
    "            seq=batch['seq'][i],  # num_res x 1\n",
    "            coords=batch['coords'][i],  # num_res x 3 x 3\n",
    "            node_s=batch['node_s'][i],  # num_res x num_conf x 4\n",
    "            node_v=batch['node_v'][i],  # num_res x num_conf x 4 x 3\n",
    "            edge_s=batch['edge_s'][i],  # num_edges x num_conf x 32\n",
    "            edge_v=batch['edge_v'][i],  # num_edges x num_conf x 1 x 3\n",
    "            edge_index=batch['edge_index'][i],  # 2 x num_edges\n",
    "            mask=batch['mask'][i]  # num_res x 1\n",
    "        ))\n",
    "\n",
    "    return Batch.from_data_list(data_list), batch_size, length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct_data = dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['seq', 'coords', 'node_s', 'node_v', 'edge_s', 'edge_v', 'edge_index', 'mask', 'z_t'])"
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import tree\n",
    "# 所有变量都增加一个batch维度\n",
    "struct_data = tree.map_structure(lambda x:\n",
    "                                x.unsqueeze(0).repeat_interleave(1, dim=0),\n",
    "                                struct_data)\n",
    "struct_data.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 76, 4])\n"
     ]
    }
   ],
   "source": [
    "print(struct_data['node_s'].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2 torch.Size([76, 4])\n",
      "1 torch.Size([76])\n",
      "3 torch.Size([76, 3, 3])\n",
      "2 torch.Size([76, 4])\n",
      "3 torch.Size([76, 4, 3])\n",
      "2 torch.Size([760, 32])\n",
      "3 torch.Size([760, 1, 3])\n",
      "2 torch.Size([2, 760])\n",
      "1 torch.Size([76])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3680641/580161542.py:22: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. To learn more, see the migration guide https://numpy.org/devdocs/numpy_2_0_migration_guide.html#adapting-to-changes-in-the-copy-keyword\n",
      "  batch_data[key][i,:lengths[i]] = samples[i][key]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "could not broadcast input array from shape (760,32) into shape (76,32)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[119], line 24\u001b[0m\n\u001b[1;32m     22\u001b[0m             batch_data[key][i,:lengths[i]] \u001b[38;5;241m=\u001b[39m samples[i][key]\n\u001b[1;32m     23\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m batch_data\n\u001b[0;32m---> 24\u001b[0m batch_data \u001b[38;5;241m=\u001b[39m \u001b[43mcollate_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdataset\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataset\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataset\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataset\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[119], line 22\u001b[0m, in \u001b[0;36mcollate_fn\u001b[0;34m(samples)\u001b[0m\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(batch_size):\n\u001b[1;32m     21\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m key \u001b[38;5;129;01min\u001b[39;00m keys:\n\u001b[0;32m---> 22\u001b[0m         \u001b[43mbatch_data\u001b[49m\u001b[43m[\u001b[49m\u001b[43mkey\u001b[49m\u001b[43m]\u001b[49m\u001b[43m[\u001b[49m\u001b[43mi\u001b[49m\u001b[43m,\u001b[49m\u001b[43m:\u001b[49m\u001b[43mlengths\u001b[49m\u001b[43m[\u001b[49m\u001b[43mi\u001b[49m\u001b[43m]\u001b[49m\u001b[43m]\u001b[49m \u001b[38;5;241m=\u001b[39m samples[i][key]\n\u001b[1;32m     23\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m batch_data\n",
      "\u001b[0;31mValueError\u001b[0m: could not broadcast input array from shape (760,32) into shape (76,32)"
     ]
    }
   ],
   "source": [
    "def collate_fn(samples):\n",
    "    batch_size = len(samples)\n",
    "    lengths = [s[\"coords\"].shape[0] for s in samples]  # 获取RNA序列长度\n",
    "    max_len = max(lengths)  # 当前batch最大长度\n",
    "    keys = ['z_t', 'seq', 'coords', 'node_s', 'node_v', 'edge_s', 'edge_v', 'edge_index', 'mask']\n",
    "    batch_data = {}\n",
    "    for key in keys:\n",
    "        shape = samples[0][key].shape\n",
    "        print(len(shape), shape)\n",
    "        if key == 'edge_index':\n",
    "            value = np.zeros((batch_size, shape[0], shape[1]), dtype=np.float32)\n",
    "        elif len(shape) == 1:\n",
    "            value = np.zeros((batch_size, max_len), dtype=np.float32)\n",
    "        elif len(shape) == 2:\n",
    "            value = np.zeros((batch_size, max_len, shape[1]), dtype=np.float32)\n",
    "        elif len(shape) == 3:\n",
    "            value = np.zeros((batch_size, max_len, shape[1], shape[2]), dtype=np.float32)\n",
    "        batch_data[key] = value\n",
    "\n",
    "    for i in range(batch_size):\n",
    "        for key in keys:\n",
    "            batch_data[key][i,:lengths[i]] = samples[i][key]\n",
    "    return batch_data\n",
    "batch_data = collate_fn((dataset[0], dataset[1], dataset[2], dataset[3]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1668, 4])"
      ]
     },
     "execution_count": 113,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_data['z_t'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch, batch_size, length = geo_batch(struct_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_t = batch.z_t\n",
    "cond_x = torch.zeros_like(z_t)\n",
    "\n",
    "init_seq = torch.cat([z_t, cond_x], -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([76, 8])"
      ]
     },
     "execution_count": 86,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "init_seq.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([76, 4])"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z_t.shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pyg",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
